{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# 查询SQL数据库\n",
    "加载 SQL 数据库"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "ea7c3c9d09feddbb"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "from langchain_community.utilities import SQLDatabase\n",
    "\n",
    "load_dotenv()\n",
    "db = SQLDatabase.from_uri(os.getenv(\"SQLlITE_DATABASE_URL\"))\n",
    "\n",
    "\n",
    "def get_schema(_):\n",
    "    return db.get_table_info()\n",
    "\n",
    "\n",
    "def run_query(query):\n",
    "    return db.run(query)\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-20T05:23:36.369543Z",
     "start_time": "2024-11-20T05:23:36.308057Z"
    }
   },
   "id": "f2fcc0f934bb4ff4",
   "execution_count": 2
  },
  {
   "cell_type": "markdown",
   "source": [],
   "metadata": {
    "collapsed": false
   },
   "id": "97c2bd93da6c9231"
  },
  {
   "cell_type": "markdown",
   "source": [],
   "metadata": {
    "collapsed": false
   },
   "id": "b7f63feb470a6da9"
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "'要查询有多少员工，可以使用`COUNT`函数来统计`Employee`表中的记录数。下面是相应的SQL查询语句：\\n\\n```sql\\nSELECT COUNT(*) AS NumberOfEmployees \\nFROM Employee;\\n```\\n\\n这条SQL语句会返回一个结果集，其中包含一列`NumberOfEmployees`，显示了`Employee`表中所有记录的数量，即员工总数。'"
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "template = \"\"\"根据下面的表模式，仅编写段SQL查询来回答用户的问题：\n",
    "{schema}\n",
    "\n",
    "问题：{question}\n",
    "SQL查询：\n",
    "\"\"\"\n",
    "llm = ChatOpenAI(\n",
    "    # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key=\"sk-xxx\",\n",
    "    openai_api_key=os.getenv(\"DASHSCOPE_API_KEY\"),\n",
    "    openai_api_base=\"https://dashscope.aliyuncs.com/compatible-mode/v1\",\n",
    "    model_name=\"qwen-max\",\n",
    "    verbose=True\n",
    ")\n",
    "prompt = ChatPromptTemplate.from_template(template)\n",
    "\n",
    "\n",
    "def print_sql(query):\n",
    "    print(query)\n",
    "    return query\n",
    "\n",
    "\n",
    "sql_chain = RunnablePassthrough.assign(schema=get_schema) | prompt | llm.bind(\n",
    "    stop=[\"\\nSQL语句:\"]) | StrOutputParser()\n",
    "sql_chain.invoke({\"question\": \"有多少员工？\"})\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-20T05:23:41.996199Z",
     "start_time": "2024-11-20T05:23:36.371953Z"
    }
   },
   "id": "b911f624a28f8ef5",
   "execution_count": 3
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 使用链来实现\n",
    "链支持那些步骤明确且可预测的应用场景。在这个场景下，我们可以创建一个简单的链，来执行以下操作：\n",
    "\n",
    "1. 将问题转换为SQL查询：我们可以利用LangChain的create_sql_query_chain方法将用户问题转换为SQL查询。\n",
    "2. 执行查询：通过SQL查询数据库获取结果。\n",
    "3. 根据结果回答问题：生成用户问题的最终回答。\n",
    "\n",
    "将问题转换为SQL查询"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "f1ad9157d3084722"
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "'[(8,)]'"
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "from langchain.chains.sql_database.query import create_sql_query_chain\n",
    "\n",
    "# 将问题转换为SQL查询\n",
    "sql_query_chain = create_sql_query_chain(llm, db)\n",
    "sql_response = sql_query_chain.invoke({\"question\": \"有多少员工？\"})\n",
    "db.run(sql_response[9:])"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-20T05:23:44.224124Z",
     "start_time": "2024-11-20T05:23:41.997997Z"
    }
   },
   "id": "be0d900bdd34300",
   "execution_count": 4
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 使用链(chains)实现问答\n",
    "接下来我们将构建一个链来处理SQL查询并返回结果。该链的步骤如下：\n",
    "\n",
    "1. 将用户问题转换为SQL查询\n",
    "2. 执行查询\n",
    "3. 使用查询结果回答问题\n",
    "\n",
    "## 创建并执行SQL查询链\n",
    "我们使用LangChain提供的`create_sql_query_chain`来将用户问题转换为SQL查询。然后使用QuerySQLDataBaseTool执行查询并通过链（Chain）将查询结果返回给用户。\n",
    "\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "691c41d026bb0f47"
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'question': '生日是1970年前的员工', 'query': '```sql\\nSELECT \"LastName\", \"FirstName\", \"BirthDate\" \\nFROM \"Employee\" \\nWHERE \"BirthDate\" < \\'1970-01-01\\' \\nLIMIT 5;\\n```'}\n",
      "```sql\n",
      "SELECT \"LastName\", \"FirstName\", \"BirthDate\" \n",
      "FROM \"Employee\" \n",
      "WHERE \"BirthDate\" < '1970-01-01' \n",
      "LIMIT 5;\n",
      "```\n"
     ]
    },
    {
     "data": {
      "text/plain": "'看起来在执行SQL查询时遇到了一个语法错误。根据提供的信息，SQL语句本身看起来是正确的，用于从`Employee`表中选择出生日期早于1970年1月1日的员工的名字、姓氏和出生日期，并限制结果为前5条记录。\\n\\n错误可能来源于几个方面：\\n1. **数据库连接问题**：可能是尝试运行该命令时数据库没有正确连接。\\n2. **表或列名错误**：如果`Employee`表不存在或者`LastName`, `FirstName`, `BirthDate`这些字段名称不正确的话也会报错。\\n3. **SQL关键字使用不当**：虽然在这个例子中不太可能是这个问题，但在其他情况下使用了保留字作为表名或字段名而未加引号可能会导致类似的问题。\\n4. **格式问题**：有时候复制粘贴代码时可能会不小心引入了不可见字符或其他格式问题。\\n\\n建议检查以下几点来解决问题：\\n- 确保你的数据库服务正在运行并且你已经成功地与之建立了连接。\\n- 核实`Employee`表确实存在于你的数据库中，并且包含了`LastName`, `FirstName`, 和 `BirthDate`这三个字段。\\n- 尝试直接在数据库管理工具（如SQLite命令行界面、phpMyAdmin等）里手动输入这条SQL语句看看是否能够正常工作。\\n- 检查是否有任何多余的空格或特殊字符被无意间包含进去了。\\n\\n如果你确定所有上述条件都满足但仍然遇到同样的错误，请提供更多关于错误的具体上下文信息以便进一步诊断。'"
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_core.prompts import PromptTemplate\n",
    "from langchain_community.tools import QuerySQLDataBaseTool\n",
    "\n",
    "template = \"\"\"Given the following user question, corresponding SQL query, and SQL result, answer the user question.\n",
    "\n",
    "Question: {question}\n",
    "SQL Query: {query}\n",
    "SQL Result: {result}\n",
    "Answer: \"\"\"\n",
    "answer_prompt = PromptTemplate.from_template(template)\n",
    "# 将问题转换为SQL查询\n",
    "write_query = create_sql_query_chain(llm, db)\n",
    "execute_query = QuerySQLDataBaseTool(db=db)\n",
    "\n",
    "\n",
    "# 清理查询语句\n",
    "def clean_query(query):\n",
    "    print(query)\n",
    "    sql_new = query.get(\"query\", \"\").replace(\"SQLQuery:\", \"\").strip()\n",
    "    print(sql_new)\n",
    "    return sql_new\n",
    "\n",
    "\n",
    "chain = RunnablePassthrough.assign(query=write_query).assign(\n",
    "    result=lambda x: execute_query.invoke({\"query\": clean_query(x)})\n",
    ") | answer_prompt | llm | StrOutputParser()\n",
    "\n",
    "chain.invoke({\"question\": \"计算有多少员工。\"})"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-20T05:25:22.943488Z",
     "start_time": "2024-11-20T05:25:01.295858Z"
    }
   },
   "id": "ecbb7910d9209b96",
   "execution_count": 9
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 使用SQLDatabaseChain实现\n"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "4ebb8a0b5ef487f7"
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001B[1m> Entering new SQLDatabaseChain chain...\u001B[0m\n",
      "计算有多少员工。\n",
      "SQLQuery:\u001B[32;1m\u001B[1;3mSQLQuery: SELECT COUNT(\"EmployeeId\") AS \"NumberOfEmployees\" FROM \"Employee\";\u001B[0m\n",
      "SQLResult: \u001B[33;1m\u001B[1;3m[(8,)]\u001B[0m\n",
      "Answer:\u001B[32;1m\u001B[1;3m有8名员工。\u001B[0m\n",
      "\u001B[1m> Finished chain.\u001B[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": "{'query': '计算有多少员工。', 'result': '有8名员工。'}"
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_experimental.sql import SQLDatabaseChain\n",
    "\n",
    "sql_database_chain = SQLDatabaseChain.from_llm(llm, db, verbose=True)\n",
    "sql_database_chain.invoke(\"计算有多少员工。\")"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-20T05:23:51.691112Z",
     "start_time": "2024-11-20T05:23:47.695899Z"
    }
   },
   "id": "e9804d2aac0f64bf",
   "execution_count": 6
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 使用代理(Agents)实现问答\n",
    "LangChain拥有一种称为SQL代理的工具，它能够比`SQLDatabaseChain`更灵活地与SQL数据库进行交互。\n",
    "\n",
    "使用SQL代理的主要优点包括：\n",
    "\n",
    "- 它可以根据数据库模式以及数据库内容（例如描述特定表）回答问题。\n",
    "- 它可以通过运行生成的查询、捕获回溯信息并重新正确生成来从错误中恢复。\n",
    "\n",
    "要初始化代理，我们使用`create_sql_agent`函数。\n",
    "\n",
    "该代理包含`SQLDatabaseToolkit`，其中包含以下工具：\n",
    "\n",
    "- 创建和执行查询\n",
    "- 检查查询语法\n",
    "- 检索表描述\n",
    "- …等等\n",
    "\n",
    "代理提供了比链更灵活的方式，允许多次查询数据库，并根据需要生成答案。\n",
    "\n",
    "## 初始化SQL代理\n",
    " 使用`create_sql_agent`函数、`SQLDatabaseToolkit`\n",
    " \n",
    "API参考：\n",
    "\n",
    "- [create_sql_agent](https://api.python.langchain.com/en/latest/agents/langchain.agents.agent_toolkits.sql.base.create_sql_agent.html)\n",
    "- [SQLDatabaseToolkit](https://api.python.langchain.com/en/latest/agents/langchain.agents.agent_toolkits.sql.toolkit.SQLDatabaseToolkit.html)\n",
    "- [AgentExecutor](https://api.python.langchain.com/en/latest/agents/langchain.agents.agent.AgentExecutor.html)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "fa8e8e0b11ea84a7"
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "name='SQL Agent Executor' verbose=True agent=RunnableAgent(runnable=RunnableAssign(mapper={\n",
      "  agent_scratchpad: RunnableLambda(lambda x: format_log_to_str(x['intermediate_steps']))\n",
      "})\n",
      "| PromptTemplate(input_variables=['agent_scratchpad', 'input'], input_types={}, partial_variables={'tools': \"sql_db_query - Input to this tool is a detailed and correct SQL query, output is a result from the database. If the query is not correct, an error message will be returned. If an error is returned, rewrite the query, check the query, and try again. If you encounter an issue with Unknown column 'xxxx' in 'field list', use sql_db_schema to query the correct table fields.\\nsql_db_schema - Input to this tool is a comma-separated list of tables, output is the schema and sample rows for those tables. Be sure that the tables actually exist by calling sql_db_list_tables first! Example Input: table1, table2, table3\\nsql_db_list_tables - Input is an empty string, output is a comma-separated list of tables in the database.\\nsql_db_query_checker - Use this tool to double check if your query is correct before executing it. Always use this tool before executing a query with sql_db_query!\", 'tool_names': 'sql_db_query, sql_db_schema, sql_db_list_tables, sql_db_query_checker'}, template='You are an agent designed to interact with a SQL database.\\nGiven an input question, create a syntactically correct sqlite query to run, then look at the results of the query and return the answer.\\nUnless the user specifies a specific number of examples they wish to obtain, always limit your query to at most 10 results.\\nYou can order the results by a relevant column to return the most interesting examples in the database.\\nNever query for all the columns from a specific table, only ask for the relevant columns given the question.\\nYou have access to tools for interacting with the database.\\nOnly use the below tools. Only use the information returned by the below tools to construct your final answer.\\nYou MUST double check your query before executing it. If you get an error while executing a query, rewrite the query and try again.\\n\\nDO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database.\\n\\nIf the question does not seem related to the database, just return \"I don\\'t know\" as the answer.\\n\\n\\n{tools}\\n\\nUse the following format:\\n\\nQuestion: the input question you must answer\\nThought: you should always think about what to do\\nAction: the action to take, should be one of [{tool_names}]\\nAction Input: the input to the action\\nObservation: the result of the action\\n... (this Thought/Action/Action Input/Observation can repeat N times)\\nThought: I now know the final answer\\nFinal Answer: the final answer to the original input question\\n\\nBegin!\\n\\nQuestion: {input}\\nThought: I should look at the tables in the database to see what I can query.  Then I should query the schema of the most relevant tables.\\n{agent_scratchpad}')\n",
      "| RunnableBinding(bound=ChatOpenAI(client=<openai.resources.chat.completions.Completions object at 0x111ad97f0>, async_client=<openai.resources.chat.completions.AsyncCompletions object at 0x1120c6c00>, root_client=<openai.OpenAI object at 0x111b36bd0>, root_async_client=<openai.AsyncOpenAI object at 0x1120c51f0>, model_name='qwen-max', model_kwargs={}, openai_api_key=SecretStr('**********'), openai_api_base='https://dashscope.aliyuncs.com/compatible-mode/v1'), kwargs={'stop': ['\\nObservation']}, config={}, config_factories=[])\n",
      "| ReActSingleInputOutputParser(), input_keys_arg=['input'], return_keys_arg=['output'], stream_runnable=True) tools=[QuerySQLDataBaseTool(description=\"Input to this tool is a detailed and correct SQL query, output is a result from the database. If the query is not correct, an error message will be returned. If an error is returned, rewrite the query, check the query, and try again. If you encounter an issue with Unknown column 'xxxx' in 'field list', use sql_db_schema to query the correct table fields.\", db=<langchain_community.utilities.sql_database.SQLDatabase object at 0x110ad3770>), InfoSQLDatabaseTool(description='Input to this tool is a comma-separated list of tables, output is the schema and sample rows for those tables. Be sure that the tables actually exist by calling sql_db_list_tables first! Example Input: table1, table2, table3', db=<langchain_community.utilities.sql_database.SQLDatabase object at 0x110ad3770>), ListSQLDatabaseTool(db=<langchain_community.utilities.sql_database.SQLDatabase object at 0x110ad3770>), QuerySQLCheckerTool(description='Use this tool to double check if your query is correct before executing it. Always use this tool before executing a query with sql_db_query!', db=<langchain_community.utilities.sql_database.SQLDatabase object at 0x110ad3770>, llm=ChatOpenAI(client=<openai.resources.chat.completions.Completions object at 0x111ad97f0>, async_client=<openai.resources.chat.completions.AsyncCompletions object at 0x1120c6c00>, root_client=<openai.OpenAI object at 0x111b36bd0>, root_async_client=<openai.AsyncOpenAI object at 0x1120c51f0>, model_name='qwen-max', model_kwargs={}, openai_api_key=SecretStr('**********'), openai_api_base='https://dashscope.aliyuncs.com/compatible-mode/v1'), llm_chain=LLMChain(verbose=False, prompt=PromptTemplate(input_variables=['dialect', 'query'], input_types={}, partial_variables={}, template='\\n{query}\\nDouble check the {dialect} query above for common mistakes, including:\\n- Using NOT IN with NULL values\\n- Using UNION when UNION ALL should have been used\\n- Using BETWEEN for exclusive ranges\\n- Data type mismatch in predicates\\n- Properly quoting identifiers\\n- Using the correct number of arguments for functions\\n- Casting to the correct data type\\n- Using the proper columns for joins\\n\\nIf there are any of the above mistakes, rewrite the query. If there are no mistakes, just reproduce the original query.\\n\\nOutput the final SQL query only.\\n\\nSQL Query: '), llm=ChatOpenAI(client=<openai.resources.chat.completions.Completions object at 0x111ad97f0>, async_client=<openai.resources.chat.completions.AsyncCompletions object at 0x1120c6c00>, root_client=<openai.OpenAI object at 0x111b36bd0>, root_async_client=<openai.AsyncOpenAI object at 0x1120c51f0>, model_name='qwen-max', model_kwargs={}, openai_api_key=SecretStr('**********'), openai_api_base='https://dashscope.aliyuncs.com/compatible-mode/v1'), output_parser=StrOutputParser(), llm_kwargs={}))]\n"
     ]
    }
   ],
   "source": [
    "\n",
    "from langchain.agents import AgentType\n",
    "from langchain_community.agent_toolkits import SQLDatabaseToolkit, create_sql_agent\n",
    "\n",
    "toolkit = SQLDatabaseToolkit(db=db, llm=llm)\n",
    "tools = toolkit.get_tools()\n",
    "\n",
    "agent_executor = create_sql_agent(llm=llm, toolkit=toolkit, verbose=True,agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,)\n",
    "print(agent_executor)\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-20T05:23:52.453760Z",
     "start_time": "2024-11-20T05:23:51.693769Z"
    }
   },
   "id": "bdf74cad0a359c21",
   "execution_count": 7
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001B[1m> Entering new SQL Agent Executor chain...\u001B[0m\n",
      "\u001B[32;1m\u001B[1;3mAction: sql_db_list_tables\n",
      "Action Input: \u001B[0m\u001B[38;5;200m\u001B[1;3mAlbum, Artist, Customer, Employee, Genre, Invoice, InvoiceLine, MediaType, Playlist, PlaylistTrack, Track\u001B[0m\u001B[32;1m\u001B[1;3m看起来数据库中的表都与音乐相关，如专辑、艺术家、客户等信息，并没有关于天气的数据。因此，对于“今天天气”这样的问题，基于当前可用的数据库信息是无法回答的。\n",
      "\n",
      "Final Answer: I don't know\u001B[0m\n",
      "\n",
      "\u001B[1m> Finished chain.\u001B[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": "\"I don't know\""
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_executor.run(\"我需要10人生日是1970年之前的员工\")"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-11-20T07:02:36.332554Z",
     "start_time": "2024-11-20T07:02:30.863825Z"
    }
   },
   "id": "d9a3078da35835cd",
   "execution_count": 11
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
